open Transfer

(* 这里假设调用约定：
   - 函数参数通过栈传递
   - 返回值放在 a0
   - 保存 s0(fp), ra 寄存器
   - 临时变量用 t0-t6 寄存器循环使用
*)

let max_tmp_regs = 7  (* t0 ~ t6 *)

(* 临时变量映射到寄存器：t0~t6循环使用 *)
let reg_of_temp temp =
  try
    Scanf.sscanf temp "t%d" (fun x -> Printf.sprintf "t%d" (x mod max_tmp_regs))
  with _ -> failwith ("Invalid temp: " ^ temp)

(* 栈上偏移管理 *)
module VarStackMap = struct
  (* 变量名 -> 栈偏移 *)
  let tbl = Hashtbl.create 32
  let next_offset = ref 0

  let clear () =
    Hashtbl.clear tbl;
    next_offset := 0

  let add var =
    if not (Hashtbl.mem tbl var) then begin
      next_offset := !next_offset + 4;  (* 每个变量4字节 *)
      Hashtbl.add tbl var (- !next_offset)
    end

  let get var =
    if Hashtbl.mem tbl var then Hashtbl.find tbl var else failwith ("Variable not found on stack: " ^ var)

  let stack_size () = !next_offset
end

(* 辅助指令生成 *)
let emit li =
  Printf.sprintf "%s" li

let load_word rd offset base =
  Printf.sprintf "lw %s, %d(%s)" rd offset base

let store_word rs offset base =
  Printf.sprintf "sw %s, %d(%s)" rs offset base

let addi rd rs imm =
  Printf.sprintf "addi %s, %s, %d" rd rs imm

let mv rd rs =
  Printf.sprintf "mv %s, %s" rd rs

let li rd imm =
  Printf.sprintf "li %s, %s" rd imm

let label l = l ^ ":"

let j l = "j " ^ l

let jalr rs = Printf.sprintf "jalr %s" rs

let beqz rs l = Printf.sprintf "beqz %s, %s" rs l

let bnez rs l = Printf.sprintf "bnez %s, %s" rs l

let neg rd rs = Printf.sprintf "neg %s, %s" rd rs

(* 二元操作映射 *)
let binop_to_riscv op =
  match op with
  | "+" -> "add"
  | "-" -> "sub"
  | "*" -> "mul"
  | "/" -> "div"
  | "%" -> "rem"
  | "&&" -> "and"
  | "||" -> "or"
  | _ -> failwith ("Unsupported binop for codegen: " ^ op)

(* 代码生成主函数 *)
let gen_riscv_code tac_list =
  let code = ref [] in

  (* 记录当前函数名称 *)
  let current_func = ref "" in

  (* 参数压栈顺序，简单模拟 *)
  let param_stack = ref [] in

  (* emit一条指令 *)
  let emit s = code := !code @ [s] in

  (* 把临时变量的值加载到寄存器 *)
  let load_var var reg =
    (* 变量名有两种：tX 临时变量，或者形如 x_1_func 这种SSA变量 *)
    if String.length var > 0 && var.[0] = 't' then
      (* 临时变量直接映射寄存器 *)
      let r = reg_of_temp var in
      if r <> reg then emit (mv reg r)
    else
      (* 栈变量，先加载到 reg *)
      let off = VarStackMap.get var in
      emit (load_word reg off "s0")
  in

  (* 把寄存器的值存回变量 *)
  let store_var var reg =
    if String.length var > 0 && var.[0] = 't' then
      let r = reg_of_temp var in
      if r <> reg then emit (mv r reg)
    else
      let off = VarStackMap.get var in
      emit (store_word reg off "s0")
  in

  (* 生成函数入口代码，分配栈帧 *)
  let gen_func_prologue fname =
    current_func := fname;
    VarStackMap.clear ();
    param_stack := [];
    emit (label fname);
    (* 保存 ra 和 fp *)
    emit (store_word "ra" (-4) "sp");
    emit (store_word "s0" (-8) "sp");
    (* 设置新的帧指针 *)
    emit (addi "s0" "sp" (-8));
    (* 预留栈空间，暂时 0，后面根据变量分配更新 *)
    emit (addi "sp" "sp" 0)
  in

  (* 生成函数尾声，释放栈帧 *)
  let gen_func_epilogue () =
    let _ = VarStackMap.stack_size () in
    emit (addi "sp" "s0" 0);  (* 恢复 sp = s0 *)
    emit (load_word "ra" (-4) "sp");
    emit (load_word "s0" (-8) "sp");
    emit "ret"
  in

  (* 生成一条 TAC 指令 *)
  let gen_one tac =
    match tac with
    | TacComment (_, s, _) -> emit ("# " ^ s)

    | TacLabel l -> emit (label l)

    | TacGoto l -> emit (j l)

    | TacAssign (dst, src) ->
        (* dst和src处理：先确认变量分配栈偏移 *)
        VarStackMap.add dst;
        if String.for_all (fun c -> c >= '0' && c <= '9') src then
          (* 立即数 *)
          let reg = "t0" in
          emit (li reg src);
          store_var dst reg
        else begin
          VarStackMap.add src;
          let reg = "t0" in
          load_var src reg;
          store_var dst reg
        end

    | TacBinOp (dst, lhs, op, rhs) ->
        VarStackMap.add dst;
        VarStackMap.add lhs;
        VarStackMap.add rhs;
        let rd = "t0" in
        let rs1 = "t1" in
        let rs2 = "t2" in
        load_var lhs rs1;
        load_var rhs rs2;
        (match op with
        | "+" | "-" | "*" | "/" | "%" | "&&" | "||" ->
            let inst = binop_to_riscv op in
            emit (Printf.sprintf "%s %s, %s, %s" inst rd rs1 rs2)
        | "==" | "!=" | "<" | "<=" | ">" | ">=" ->
            (* 处理比较，使用sub + seqz/snez 或 slt *)
            (match op with
            | "==" ->
                emit (Printf.sprintf "sub %s, %s, %s" rd rs1 rs2);
                emit (Printf.sprintf "seqz %s, %s" rd rd)
            | "!=" ->
                emit (Printf.sprintf "sub %s, %s, %s" rd rs1 rs2);
                emit (Printf.sprintf "snez %s, %s" rd rd)
            | "<" ->
                emit (Printf.sprintf "slt %s, %s, %s" rd rs1 rs2)
            | "<=" ->
                emit (Printf.sprintf "slt %s, %s, %s" rd rs2 rs1);
                emit (Printf.sprintf "seqz %s, %s" rd rd)
            | ">" ->
                emit (Printf.sprintf "slt %s, %s, %s" rd rs2 rs1)
            | ">=" ->
                emit (Printf.sprintf "slt %s, %s, %s" rd rs1 rs2);
                emit (Printf.sprintf "seqz %s, %s" rd rd)
            | _ -> ())
        | _ -> failwith ("Unsupported binop: " ^ op));
        store_var dst rd

    | TacUnOp (dst, op, src) ->
        VarStackMap.add dst;
        VarStackMap.add src;
        let rd = "t0" in
        let rs = "t1" in
        load_var src rs;
        (match op with
        | "+" -> emit (mv rd rs)
        | "-" -> emit (neg rd rs)
        | "!" -> emit (Printf.sprintf "seqz %s, %s" rd rs)
        | _ -> failwith ("Unsupported unop: " ^ op));
        store_var dst rd

    | TacIfGoto (cond, label) ->
        VarStackMap.add cond;
        let rc = "t0" in
        load_var cond rc;
        emit (bnez rc label)

    | TacParam t ->
        VarStackMap.add t;
        (* 参数压栈，先load寄存器，再存到栈上 *)
        let rt = "t0" in
        load_var t rt;
        (* 按调用顺序压栈 *)
        let offset = 4 * (List.length !param_stack) + 8 in
        emit (store_word rt offset "s0");
        param_stack := !param_stack @ [offset]

    | TacCall { dst; fname; args } ->
    let argc = List.length args in
    (* ↓↓↓ 下面保持和原来逻辑一致，只要把 argc 换成上面的局部变量 ↓↓↓ *)
    let param_space = 4 * argc in
    emit (Printf.sprintf "addi sp, sp, -%d" param_space);
    emit (Printf.sprintf "call %s" fname);
    emit (Printf.sprintf "addi sp, sp, %d" param_space);
    VarStackMap.add dst;
    store_var dst "a0";
    param_stack := []

    | TacReturn None ->
        gen_func_epilogue ()

    | TacReturn (Some t) ->
        VarStackMap.add t;
        let rt = "t0" in
        load_var t rt;
        emit (mv "a0" rt);
        gen_func_epilogue ()

    | TacPhi (dst, v1, v2) ->
        (* phi 在寄存器分配阶段消除，先简单生成注释 *)
        emit (Printf.sprintf "# phi %s = phi(%s, %s) (handled by register allocator)" dst v1 v2)
  in

  (* 按函数划分，先找标签为函数名的Tac *)
  let is_func_label l =
  let len = String.length l in
  if len = 0 then false
  else
    let c = l.[0] in
    c >= 'a' && c <= 'z' &&
    let rec check i =
      if i >= len then true
      else
        let c = l.[i] in
        if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') then check (i+1)
        else false
    in
    check 1

  in
  let rec split_funcs tac_list acc current =
  match tac_list with
  | [] -> List.rev (List.rev current :: acc)
  | TacLabel l :: rest when is_func_label l ->
      if current = [] then
        split_funcs rest acc [TacLabel l]
      else
        split_funcs rest (List.rev current :: acc) [TacLabel l]
  | hd :: rest -> split_funcs rest acc (hd :: current)

      in
  let funcs = split_funcs tac_list [] [] in

  (* 生成代码 *)
  List.iter (fun func_tacs ->
    match func_tacs with
    | TacLabel fname :: body ->
      gen_func_prologue fname;
      List.iter gen_one body;
      (* 结尾可能没有 return，需要补充 *)
      let last = List.nth func_tacs ((List.length func_tacs) - 1) in
      (match last with
      | TacReturn _ -> ()
      | _ -> gen_func_epilogue ())
    | _ -> ()) funcs;

  !code
